package mllib;

import Jama.EigenvalueDecomposition;
import Jama.Matrix;
import common.ConstVars;
import image.ImageFixer;
import org.apache.commons.math.linear.Array2DRowRealMatrix;
import org.apache.commons.math.linear.ArrayRealVector;
import org.apache.commons.math.linear.RealMatrix;
import org.apache.commons.math.linear.RealVector;

import java.awt.image.BufferedImage;
import java.io.*;
import java.util.Scanner;

public class PCA implements Reducer
{
    public RealMatrix wTrans;
    private RealVector sampleMean;
    private double[] eigvals;

    public static void main(String[] args)
    {
      if(args.length == 0)
      {
        System.out.println("args: <pcafile>");
        System.exit(1);
      }
      PCA x = new PCA(args[0]);
      x.makeEigenFaces();
    }

    //Constructor to get wTrans and sampleMean from saved datafile
    public PCA( String filename)
    {
        try {
          Scanner sc = new Scanner(new File(filename));
          int dim = sc.nextInt();
          int newDim = sc.nextInt();
          double[][] w = new double[newDim][dim];
          double[] sm = new double[dim];
          for(int i = 0; i < newDim; i++)
              for(int j = 0; j< dim; j++)
                  w[i][j] = sc.nextDouble();
          
          for(int i = 0; i < dim; i ++)
              sm[i] = sc.nextDouble();

          wTrans = new Array2DRowRealMatrix(w);
          sampleMean = new ArrayRealVector(sm);
        }
        catch(FileNotFoundException e) {
          e.printStackTrace();
        }

    }

    public void makeEigenFaces()
    {
      for(int i = 0;i < wTrans.getRowDimension();i++)
      {
        double[] row = wTrans.getRow(i);
        if(row.length != sampleMean.getDimension())
        {
          System.out.println("Well, shit.");
          System.exit(1);
        }
        double min = 0;
        double max = 0;
        for(int j = 0;j < row.length;j++)
        {
          if(row[j] < min)
            min = row[j];
          if(row[j] > max)
            max = row[j];
        }

        BufferedImage img = new BufferedImage(ConstVars.IMG_DIM, ConstVars.IMG_DIM, BufferedImage.TYPE_INT_RGB);
        for(int x = 0;x < ConstVars.IMG_DIM;x++) {
          for(int y = 0;y < ConstVars.IMG_DIM;y++) {
            int gray = (int) ((row[x * ConstVars.IMG_DIM + y] - min) * 255 / (max - min));
            if(gray < 0)
              gray = 0;
            if(gray > 255)
              gray = 255;
            img.setRGB(x, y, gray << 16 | gray << 8 | gray);
          }
        }
        ImageFixer.save(img, String.format("%s%03d.jpg", "eigenface", i));
      }
    }

    public void makeEigenFaces2()
    {
      for(int i = 0;i < wTrans.getRowDimension();i++)
      {
        double[] row = wTrans.getRow(i);
        if(row.length != sampleMean.getDimension())
        {
          System.out.println("Well, shit.");
          System.exit(1);
        }
        for(int j = 0;j < row.length;j++)
          row[j] += sampleMean.getEntry(j);

        BufferedImage img = new BufferedImage(ConstVars.IMG_DIM, ConstVars.IMG_DIM, BufferedImage.TYPE_INT_RGB);
        for(int x = 0;x < ConstVars.IMG_DIM;x++) {
          for(int y = 0;y < ConstVars.IMG_DIM;y++) {
            int gray = (int) row[x * ConstVars.IMG_DIM + y];
            if(gray < 0)
              gray = 0;
            if(gray > 255)
              gray = 255;
            img.setRGB(x, y, gray << 16 | gray << 8 | gray);
          }
        }
        ImageFixer.save(img, String.format("%s%03d.jpg", "eigenface", i));
      }
    }

    public double[] getEigenValues()
    {
        return eigvals;
    }
    //Prints the information necessary for this PCA to be rebuilt from a file
    public void toFile(String filename)
    {
      try {
        	BufferedWriter out = new BufferedWriter(new FileWriter(filename));
            int dim = wTrans.getColumnDimension();
            int newDim = wTrans.getRowDimension();
            out.write( dim + "\n");
            out.write(newDim + "\n");
            for( int i = 0 ; i < newDim; i++)
                for(int j = 0; j < dim; j++)
                    out.write( wTrans.getEntry(i,j) + "\n" );
            for(int i = 0; i < dim; i++)
                    out.write( sampleMean.getEntry(i) + " ");
            out.close();                    
      }
      catch(IOException e) {}

    }
    
    public PCA( TrainingSet ts, double percent ) 
    {
        //Get sample mean and center data about origin
        sampleMean = ts.centerToOrigin();

        System.out.println("CREATING COVARIANCE MATRIX...");
        //get covariance of new centered data
        RealMatrix cov = ts.covariance();

        System.out.println("STARTING DECOMPOSITION...");
        //Do decomposition to get eigenvalues and eigenvectors
        EigenvalueDecomposition decomp = new EigenvalueDecomposition( new Matrix(cov.getData()) );
        eigvals = decomp.getRealEigenvalues();
        Matrix tempwTrans = decomp.getV();
        System.out.println("DECOMPOSITION FINISHED");

        //Find number of dimensions (k) to keep
        int k = POV(eigvals, percent);

        //Get the top k eigenvectors and normalize them. The decomposition 
        //sorts the eigenvectors by increasing eigenvalues, so we want the last
        //k eigenvectors.
        tempwTrans = tempwTrans.getMatrix(tempwTrans.getRowDimension() - k,
                tempwTrans.getRowDimension() - 1,
                0,
                tempwTrans.getColumnDimension() - 1);

        wTrans = new Array2DRowRealMatrix( tempwTrans.getArray() );
        for(int i = 0;i < k;i++)
        {
           for(int j = 0;j < 1209;j++) 
                System.err.print(wTrans.getEntry(i, j) + " ");
            System.err.println();
        }
        normalizeW();
    }

    //Normalizes the eigenvectors to ensure that the projections are not distorted
    private void normalizeW()
    {
        for( int k = 0; k < wTrans.getRowDimension(); k++)
        {
            double sum = 0;
            for(int c = 0; c < wTrans.getColumnDimension(); c++)
                sum += wTrans.getEntry(k,c) * wTrans.getEntry(k,c);
            sum = Math.sqrt(sum);
            for(int c = 0; c < wTrans.getColumnDimension(); c++)
                wTrans.setEntry(k,c,    wTrans.getEntry(k,c)/sum);
        }
    }

    //Projects Vector point into the k-dimension subspace defined by wTrans
    public RealVector projectPoint( RealVector point) throws IllegalArgumentException
    {
        //System.out.print( "Point going to : ");
        //System.out.println( wTrans.operate( point.subtract(sampleMean) ).getEntry(0));
        return wTrans.operate(point.subtract(sampleMean));
    }

    //Returns k, the number of dimensions PCA reduced down to
    public int reducedDimension()
    {
        return wTrans.getRowDimension();
    }

    //Calculates the number of dimensions needed to have a proportion of 
    //variance (POV) at least "percent"
    private int  POV( double[] evalues, double percent)
    {
        double sum = 0;
        for( int i = evalues.length-1; i >= 0; i--)
            sum += evalues[i];
        
        int counter = 0;
        double partialSum = 0;
        
        while( partialSum / sum < percent)
        {
            partialSum += evalues[evalues.length - counter - 1];
            counter++;
        }

        System.err.println("k: " + counter + "");
        System.err.println("percent: " + partialSum / sum);
        return counter;
    }
}
